import re
import json
from typing import Dict, List, Set, Tuple
from dataclasses import dataclass, field
from collections import defaultdict

@dataclass
class ServiceMethod:
    service_name: str
    method_name: str
    input_type: str
    output_type: str

@dataclass
class MessageType:
    name: str
    fields: Dict[str, str] = field(default_factory=dict)
    nested_types: List[str] = field(default_factory=list)

@dataclass
class DependencyEdge:
    from_service: str
    to_service: str
    dependency_type: str  # 'data_flow', 'message_type', 'field_reference'
    details: str

class ProtobufDependencyAnalyzer:
    def __init__(self, proto_content: str):
        self.proto_content = proto_content
        self.services: Dict[str, List[ServiceMethod]] = {}
        self.messages: Dict[str, MessageType] = {}
        self.dependencies: List[DependencyEdge] = []
        
    def parse_protobuf(self):
        """解析 protobuf 文件"""
        self._extract_services()
        self._extract_messages()
        
    def _extract_services(self):
        """提取服务定义"""
        service_pattern = r'service\s+(\w+)\s*\{([^}]+)\}'
        method_pattern = r'rpc\s+(\w+)\s*\(\s*(\w+)\s*\)\s*returns\s*\(\s*(\w+)\s*\)'
        
        for service_match in re.finditer(service_pattern, self.proto_content, re.MULTILINE | re.DOTALL):
            service_name = service_match.group(1)
            service_body = service_match.group(2)
            
            methods = []
            for method_match in re.finditer(method_pattern, service_body):
                method = ServiceMethod(
                    service_name=service_name,
                    method_name=method_match.group(1),
                    input_type=method_match.group(2),
                    output_type=method_match.group(3)
                )
                methods.append(method)
            
            self.services[service_name] = methods
    
    def _extract_messages(self):
        """提取消息类型定义"""
        message_pattern = r'message\s+(\w+)\s*\{([^}]+)\}'
        field_pattern = r'(\w+)\s+(\w+)\s*=\s*\d+;'
        repeated_field_pattern = r'repeated\s+(\w+)\s+(\w+)\s*=\s*\d+;'
        
        for message_match in re.finditer(message_pattern, self.proto_content, re.MULTILINE | re.DOTALL):
            message_name = message_match.group(1)
            message_body = message_match.group(2)
            
            fields = {}
            nested_types = []
            
            # 提取普通字段
            for field_match in re.finditer(field_pattern, message_body):
                field_type = field_match.group(1)
                field_name = field_match.group(2)
                fields[field_name] = field_type
                
                # 检查是否引用了其他消息类型
                if field_type not in ['string', 'int32', 'int64', 'bool', 'double', 'float']:
                    nested_types.append(field_type)
            
            # 提取 repeated 字段
            for field_match in re.finditer(repeated_field_pattern, message_body):
                field_type = field_match.group(1)
                field_name = field_match.group(2)
                fields[field_name] = f"repeated {field_type}"
                
                if field_type not in ['string', 'int32', 'int64', 'bool', 'double', 'float']:
                    nested_types.append(field_type)
            
            self.messages[message_name] = MessageType(
                name=message_name,
                fields=fields,
                nested_types=nested_types
            )
    
    def analyze_checkout_service_dependencies(self) -> Dict[str, any]:
        """专门分析 CheckoutService 的依赖关系"""
        checkout_service = self.services.get('CheckoutService', [])
        if not checkout_service:
            return {"error": "CheckoutService not found"}
        
        analysis = {
            "service_name": "CheckoutService",
            "methods": [],
            "upstream_dependencies": [],
            "downstream_dependencies": [],
            "data_dependencies": [],
            "message_flow": []
        }
        
        for method in checkout_service:
            method_analysis = self._analyze_method_dependencies(method)
            analysis["methods"].append(method_analysis)
            
            # 合并依赖关系
            analysis["upstream_dependencies"].extend(method_analysis["required_services"])
            analysis["data_dependencies"].extend(method_analysis["data_dependencies"])
            analysis["message_flow"].extend(method_analysis["message_flow"])
        
        # 去重
        analysis["upstream_dependencies"] = list(set(analysis["upstream_dependencies"]))
        
        return analysis
    
    def _analyze_method_dependencies(self, method: ServiceMethod) -> Dict[str, any]:
        """分析单个方法的依赖关系"""
        result = {
            "method_name": method.method_name,
            "input_type": method.input_type,
            "output_type": method.output_type,
            "required_services": [],
            "data_dependencies": [],
            "message_flow": []
        }
        
        # 分析输入消息的依赖
        input_deps = self._trace_message_dependencies(method.input_type)
        result["data_dependencies"].extend(input_deps)
        
        # 分析输出消息的依赖
        output_deps = self._trace_message_dependencies(method.output_type)
        result["data_dependencies"].extend(output_deps)
        
        # 推断可能需要调用的服务
        required_services = self._infer_required_services(method.input_type)
        result["required_services"] = [call['service'] for call in required_services]
        result["service_calls"] = required_services
        
        # 构建消息流
        message_flow = self._build_message_flow(method.input_type, method.output_type)
        result["message_flow"] = message_flow
        
        return result
    
    def _trace_message_dependencies(self, message_type: str) -> List[Dict[str, str]]:
        """递归跟踪消息类型的依赖关系"""
        dependencies = []
        
        if message_type not in self.messages:
            return dependencies
        
        message = self.messages[message_type]
        
        for field_name, field_type in message.fields.items():
            # 处理 repeated 类型
            actual_type = field_type.replace("repeated ", "")
            
            if actual_type in self.messages:
                dependencies.append({
                    "message": message_type,
                    "field": field_name,
                    "depends_on": actual_type,
                    "type": "message_reference"
                })
                
                # 递归查找嵌套依赖
                nested_deps = self._trace_message_dependencies(actual_type)
                dependencies.extend(nested_deps)
        
        return dependencies
    
    def _infer_required_services(self, input_type: str) -> List[Dict[str, str]]:
        """基于输入消息类型的字段依赖分析需要调用的服务和RPC方法"""
        required_services = []
        
        # 分析消息类型中引用的其他消息类型，推断可能的服务调用
        message_deps = self._trace_message_dependencies(input_type)
        service_calls = self._infer_service_calls_from_message_deps(input_type, message_deps)
        
        return service_calls
    
    def _build_message_flow(self, input_type: str, output_type: str) -> List[Dict[str, str]]:
        """构建消息流转路径"""
        return []
        # 不存在的
        if input_type == "PlaceOrderRequest" and output_type == "PlaceOrderResponse":
            return [
                {
                    "step": 1,
                    "action": "获取用户购物车",
                    "service": "CartService",
                    "method": "GetCart",
                    "input": "GetCartRequest",
                    "output": "Cart"
                },
                {
                    "step": 2,
                    "action": "获取产品信息",
                    "service": "ProductCatalogService", 
                    "method": "GetProduct",
                    "input": "GetProductRequest",
                    "output": "Product"
                },
                {
                    "step": 3,
                    "action": "货币转换",
                    "service": "CurrencyService",
                    "method": "Convert", 
                    "input": "CurrencyConversionRequest",
                    "output": "Money"
                },
                {
                    "step": 4,
                    "action": "计算运费",
                    "service": "ShippingService",
                    "method": "GetQuote",
                    "input": "GetQuoteRequest", 
                    "output": "GetQuoteResponse"
                },
                {
                    "step": 5,
                    "action": "处理支付",
                    "service": "PaymentService",
                    "method": "Charge",
                    "input": "ChargeRequest",
                    "output": "ChargeResponse"
                },
                {
                    "step": 6,
                    "action": "安排发货",
                    "service": "ShippingService", 
                    "method": "ShipOrder",
                    "input": "ShipOrderRequest",
                    "output": "ShipOrderResponse"
                },
                {
                    "step": 7,
                    "action": "发送订单确认",
                    "service": "EmailService",
                    "method": "SendOrderConfirmation",
                    "input": "SendOrderConfirmationRequest",
                    "output": "Empty"
                }
            ]
        return []
    
    def generate_dependency_graph(self) -> Dict[str, any]:
        """生成完整的服务依赖图"""
        graph = {
            "nodes": [],
            "edges": [],
            "service_dependencies": {}
        }
        
        # 添加所有服务作为节点
        for service_name in self.services.keys():
            graph["nodes"].append({
                "id": service_name,
                "type": "service",
                "methods": [m.method_name for m in self.services[service_name]]
            })
        
        # 分析每个服务的依赖关系
        for service_name, methods in self.services.items():
            dependencies = []
            
            for method in methods:
                # 基于消息类型依赖推断依赖
                method_deps = self._infer_service_dependencies(service_name, method)
                dependencies.extend(method_deps)
            
            graph["service_dependencies"][service_name] = dependencies
            
            # 添加边
            for dep in dependencies:
                if 'service' in dep:  # 对于 service call 类型的依赖
                    graph["edges"].append({
                        "from": service_name,
                        "to": dep['service'],
                        "type": "rpc_call",
                        "method": dep['method'],
                        "reason": dep['reason']
                    })
        
        return graph
    
    def _infer_service_dependencies(self, service_name: str, method: ServiceMethod) -> List[Dict[str, str]]:
        """基于消息类型依赖推断服务依赖关系"""
        dependencies = []
        
        # 分析输入消息类型的依赖
        input_deps = self._infer_service_calls_from_message_deps(method.input_type, 
                                                                self._trace_message_dependencies(method.input_type))
        dependencies.extend(input_deps)
        
        # 分析输出消息类型可能需要的依赖
        output_deps = self._infer_service_calls_from_message_deps(method.output_type,
                                                                 self._trace_message_dependencies(method.output_type))
        dependencies.extend(output_deps)
        
        return dependencies
    
    def _infer_service_calls_from_message_deps(self, message_type: str, message_deps: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """基于消息依赖关系推断具体的服务调用和RPC方法"""
        service_calls = []
        
        # 对于每个消息依赖，查找哪个服务返回该消息类型
        for dep in message_deps:
            depends_on_type = dep['depends_on']
            
            # 查找哪个服务的哪个方法返回这个消息类型
            for service_name, methods in self.services.items():
                for method in methods:
                    if method.output_type == depends_on_type:
                        service_calls.append({
                            'service': service_name,
                            'method': method.method_name,
                            'input_type': method.input_type,
                            'output_type': method.output_type,
                            'reason': f'需要{depends_on_type}类型的数据，可能需要调用{service_name}.{method.method_name}',
                            'dependency_path': f'{message_type}.{dep["field"]} -> {depends_on_type}'
                        })
        
        # 针对特定消息类型的特殊处理
        if message_type == 'PlaceOrderRequest':
            # PlaceOrderRequest 包含 user_id，可能需要获取购物车
            service_calls.append({
                'service': 'CartService',
                'method': 'GetCart',
                'input_type': 'GetCartRequest',
                'output_type': 'Cart',
                'reason': 'PlaceOrderRequest包含user_id，需要获取用户购物车',
                'dependency_path': 'PlaceOrderRequest.user_id -> Cart'
            })
            
            # 包含 Address，可能需要计算运费
            service_calls.append({
                'service': 'ShippingService',
                'method': 'GetQuote',
                'input_type': 'GetQuoteRequest',
                'output_type': 'GetQuoteResponse',
                'reason': 'PlaceOrderRequest包含Address，需要计算运费',
                'dependency_path': 'PlaceOrderRequest.address -> GetQuoteResponse'
            })
            
            # 包含 CreditCardInfo，需要处理支付
            service_calls.append({
                'service': 'PaymentService',
                'method': 'Charge',
                'input_type': 'ChargeRequest',
                'output_type': 'ChargeResponse',
                'reason': 'PlaceOrderRequest包含CreditCardInfo，需要处理支付',
                'dependency_path': 'PlaceOrderRequest.credit_card -> ChargeResponse'
            })
            
            # 包含 user_currency，可能需要货币转换
            service_calls.append({
                'service': 'CurrencyService',
                'method': 'Convert',
                'input_type': 'CurrencyConversionRequest',
                'output_type': 'Money',
                'reason': 'PlaceOrderRequest包含user_currency，可能需要货币转换',
                'dependency_path': 'PlaceOrderRequest.user_currency -> Money'
            })
        
        # 对于CartItem类型，可能需要获取产品信息
        if message_type == 'Cart' or any(dep['depends_on'] == 'CartItem' for dep in message_deps):
            service_calls.append({
                'service': 'ProductCatalogService',
                'method': 'GetProduct',
                'input_type': 'GetProductRequest',
                'output_type': 'Product',
                'reason': '包含CartItem，需要获取产品详细信息',
                'dependency_path': 'CartItem.product_id -> Product'
            })
        
        return service_calls
    
    def _analyze_cross_service_message_dependencies(self) -> Dict[str, List[Dict[str, str]]]:
        """分析跨服务的消息依赖关系"""
        cross_service_deps = {}
        
        for service_name, methods in self.services.items():
            service_deps = []
            
            for method in methods:
                # 分析输入消息的依赖
                input_deps = self._trace_message_dependencies(method.input_type)
                for dep in input_deps:
                    # 查找哪个服务提供这个消息类型
                    provider_services = self._find_message_providers(dep['depends_on'])
                    for provider in provider_services:
                        if provider['service'] != service_name:  # 跨服务依赖
                            service_deps.append({
                                'current_method': f'{service_name}.{method.method_name}',
                                'depends_on_service': provider['service'],
                                'depends_on_method': provider['method'],
                                'message_type': dep['depends_on'],
                                'dependency_reason': f'{method.method_name}的输入{method.input_type}依赖于{dep["depends_on"]}'
                            })
                
                # 分析输出消息可能需要的依赖（为了构造输出消息）
                output_deps = self._trace_message_dependencies(method.output_type)
                for dep in output_deps:
                    provider_services = self._find_message_providers(dep['depends_on'])
                    for provider in provider_services:
                        if provider['service'] != service_name:
                            service_deps.append({
                                'current_method': f'{service_name}.{method.method_name}',
                                'depends_on_service': provider['service'],
                                'depends_on_method': provider['method'],
                                'message_type': dep['depends_on'],
                                'dependency_reason': f'{method.method_name}的输出{method.output_type}需要构造{dep["depends_on"]}'
                            })
            
            if service_deps:
                cross_service_deps[service_name] = service_deps
        
        return cross_service_deps
    
    def _find_message_providers(self, message_type: str) -> List[Dict[str, str]]:
        """查找提供特定消息类型的服务和方法"""
        providers = []
        
        for service_name, methods in self.services.items():
            for method in methods:
                if method.output_type == message_type:
                    providers.append({
                        'service': service_name,
                        'method': method.method_name,
                        'input_type': method.input_type,
                        'output_type': method.output_type
                    })
        
        return providers
    
    def generate_mermaid_diagram(self) -> str:
        """生成 Mermaid 依赖图"""
        graph = self.generate_dependency_graph()
        
        mermaid = "graph TB\n"
        
        # 添加节点
        for node in graph["nodes"]:
            mermaid += f"    {node['id']}[{node['id']}]\n"
        
        # 添加边
        for edge in graph["edges"]:
            mermaid += f"    {edge['from']} --> {edge['to']}\n"
        
        return mermaid

# 使用示例
def analyze_otel_demo_dependencies():
    """分析 OTEL Demo 的服务依赖关系"""
    
    # 这里应该读取实际的 proto 文件内容
    with open('./data/otel-demo.proto', 'r') as f:
        proto_content = f.read()
    
    analyzer = ProtobufDependencyAnalyzer(proto_content)
    analyzer.parse_protobuf()
    
    # 专门分析 CheckoutService
    checkout_analysis = analyzer.analyze_checkout_service_dependencies()
    
    # 生成完整的依赖图
    dependency_graph = analyzer.generate_dependency_graph()
    
    # 分析跨服务消息依赖
    cross_service_deps = analyzer._analyze_cross_service_message_dependencies()
    
    # 生成 Mermaid 图
    mermaid_diagram = analyzer.generate_mermaid_diagram()
    
    return {
        "checkout_service_analysis": checkout_analysis,
        "full_dependency_graph": dependency_graph,
        "cross_service_dependencies": cross_service_deps,
        "mermaid_diagram": mermaid_diagram
    }

if __name__ == "__main__":
    results = analyze_otel_demo_dependencies()
    
    print("=== CheckoutService 依赖分析 ===")
    print(json.dumps(results["checkout_service_analysis"], indent=2, ensure_ascii=False))
    
    print("\n=== 完整服务依赖图 ===")
    print(json.dumps(results["full_dependency_graph"], indent=2, ensure_ascii=False))
    
    print("\n=== 跨服务消息依赖关系 ===")
    print(json.dumps(results["cross_service_dependencies"], indent=2, ensure_ascii=False))
    
    print("\n=== Mermaid 依赖图 ===")
    print(results["mermaid_diagram"])